{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 第一步：拉大神的项目，准备环境\n",
    "\n",
    "备注：我们的测试环境是colab的a100GPU环境，所用参数大概需要消耗37GB显存，如果显存不够可以自己改一下超参数，等下会说"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/ssbuild/chatglm_finetuning"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "把代码拷贝到工作目录"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 拷贝到colab目录，如果是自己的环境就拷贝到自己的目录\n",
    "cp -r /content/chatglm_finetuning/* /content/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 安装依赖包\n",
    "!pip install -U deep_training cpm_kernels icetk transformers>=4.26.1 deepspeed"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 第二部：制作数据集\n",
    "使用我们的merge.py 选择对应的转换文件夹，脚本会自动将文件夹下面的所有alpaca格式的json文件转换成chatglm_finetuning所需的特殊json格式，输出为output.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m无法启动 Jupyter。尝试查找 Jupyter 时出错: 运行单元格需要 jupyter 与 notebook 包。\n",
      "\u001b[1;31mRun the following command to install 'jupyter 与 notebook' into the Python environment. \n",
      "\u001b[1;31mCommand: 'python -m pip install jupyter notebook -U\n",
      "\u001b[1;31mor\n",
      "\u001b[1;31mconda install jupyter notebook -U'\n",
      "\u001b[1;31m有关详细信息，请单击 <a href=\"https://aka.ms/installJupyterForVSCode\">此处</a>。"
     ]
    }
   ],
   "source": [
    "!python merge.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 把数据集挪到data里面\n",
    "mv ./output.json ./data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 修改train_util.py的超参数\n",
    "我们的参数如下：请根据自己实际情况自行改动\n",
    "```\n",
    "train_info_args = {\n",
    "    'devices': 1,\n",
    "    'data_backend': 'record',\n",
    "    'model_type': 'chatglm',\n",
    "    # 预训练模型路径 , 因为用的是colab，直接从hunggingface拉是很快的\n",
    "    'model_name_or_path': 'THUDM/chatglm-6b',\n",
    "    'config_name': './config/config.json',\n",
    "    'tokenizer_name': 'THUDM/chatglm-6b',\n",
    "    'convert_onnx': False, # 转换onnx模型\n",
    "    'do_train': True,\n",
    "    'train_file':  [ './data/output.json'],\n",
    "    'max_epochs': 38,\n",
    "    'max_steps': -1,\n",
    "    'optimizer': 'lion', # one of adamw,adam,lamb,lion\n",
    "    'scheduler_type': 'linear',\n",
    "    'optimizer_betas': (0.9, 0.999),\n",
    "    'train_batch_size': 4,\n",
    "    'eval_batch_size': 2,\n",
    "    'test_batch_size': 2,\n",
    "    'learning_rate': 2e-5,  #\n",
    "    'adam_epsilon': 1e-8,\n",
    "    'gradient_accumulation_steps': 1,\n",
    "    'max_grad_norm': 1.0,\n",
    "    'weight_decay': 0,\n",
    "    'warmup_steps': 0,\n",
    "    'output_dir': './output',\n",
    "    'max_seq_length': 2048, # 如果资源充足，推荐长度2048 与官方保持一致\n",
    "    'max_target_length': 100,  # 预测最大长度, 保留字段\n",
    "    'use_fast_tokenizer': False,\n",
    "    'do_lower_case': False,\n",
    "\n",
    "    ##############  lora模块\n",
    "    'with_lora': True,  # 是否启用lora模块，这里必须要启动，不然小数据集下根本无法训练出好的效果\n",
    "    'inference_mode': False, # 推理模型, 不需要手动设置\n",
    "    'r': 32, # 有多好的卡，就设置多大把\n",
    "    'target_modules': ['dense','dense_h_to_4h','dense_4h_to_h','query_key_value'], # 尽量能lora的层都lora一下\n",
    "    'target_dtype': '16',\n",
    "    'lora_alpha': 32,\n",
    "    # 'enable_lora': [True],\n",
    "    'enable_lora': None,\n",
    "    'lora_dropout': 0.1,\n",
    "    'bias': 'none',  # Bias type for Lora. Can be 'none', 'all' or 'lora_only'\"\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 开启数据预处理，请记得如果数据集有更换请删除临时文件\n",
    "!python data_utils.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 模型训练，我研究过N多方案。包括自己写，能收敛有效果的只有这个\n",
    "!python train.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 测试一下吧！\n",
    "这里推荐直接在notebook加载模型，因为lora加载慢得抠脚"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from deep_training.data_helper import ModelArguments, TrainingArguments, DataArguments\n",
    "from deep_training.nlp.models.chatglm import TransformerChatGlmLMHeadModel, setup_model_profile, ChatGLMConfig,ChatGLMForConditionalGeneration\n",
    "from deep_training.nlp.models.lora import LoraArguments, LoraModel\n",
    "from transformers import HfArgumentParser\n",
    "\n",
    "from data_utils import train_info_args, NN_DataHelper,get_deepspeed_config\n",
    "from tokenization_chatglm import ChatGLMTokenizer\n",
    "\n",
    "\n",
    "class MyTransformer(TransformerChatGlmLMHeadModel, with_pl=True):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        lora_args: LoraArguments = kwargs.pop('lora_args')\n",
    "        super(MyTransformer, self).__init__(*args, **kwargs)\n",
    "        self.lora_args = lora_args\n",
    "        if lora_args.with_lora:\n",
    "            model = LoraModel(self.backbone, lora_args)\n",
    "            print('*' * 30,'lora info')\n",
    "            model.print_trainable_parameters()\n",
    "            self.set_model(model, copy_attr=False)\n",
    "\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    train_info_args['seed'] = None\n",
    "    parser = HfArgumentParser((ModelArguments, TrainingArguments, DataArguments, LoraArguments))\n",
    "    model_args, training_args, data_args, _ = parser.parse_dict(train_info_args)\n",
    "\n",
    "    setup_model_profile()\n",
    "\n",
    "    dataHelper = NN_DataHelper(model_args, training_args, data_args)\n",
    "    tokenizer: ChatGLMTokenizer\n",
    "    tokenizer, _, _, _ = dataHelper.load_tokenizer_and_config(\n",
    "        tokenizer_class_name=ChatGLMTokenizer, config_class_name=ChatGLMConfig)\n",
    "\n",
    "\n",
    "    config = ChatGLMConfig.from_pretrained('./best_ckpt')\n",
    "    config.initializer_weight = False\n",
    "\n",
    "    lora_args = LoraArguments.from_pretrained('./best_ckpt')\n",
    "\n",
    "    assert lora_args.inference_mode == True\n",
    "\n",
    "    model = MyTransformer(config=config, model_args=model_args, training_args=training_args,lora_args=lora_args)\n",
    "    # 加载lora权重\n",
    "    model.backbone.from_pretrained(model.backbone.model, pretrained_model_name_or_path = './best_ckpt', lora_config = lora_args)\n",
    "\n",
    "    base_model: ChatGLMForConditionalGeneration = model.backbone.model.model\n",
    "    # 按需修改\n",
    "    base_model.half().cuda()\n",
    "    base_model = base_model.eval()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 实验结论\n",
    "1、loss降得很低是ok的，基本不会影响模型原来的能力\n",
    "\n",
    "2、数据集小epoch可以设置得多一些，也不会影响模型能力\n",
    "\n",
    "3、如果不使用lora，则基本等于从0开始训练模型了"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据集内问题完美回答\n",
    "response, history = base_model.chat(tokenizer, \"什么是精益?\", history=[],max_length=512,\n",
    "                                        eos_token_id=config.eos_token_id,\n",
    "                                        do_sample=True, top_p=0.7, temperature=0.95,)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 写方案能力和结构化能力突然变强\n",
    "response, history = base_model.chat(tokenizer, \"帮我写一个团支部的共建方案\", history=[],max_length=512,\n",
    "                                        eos_token_id=config.eos_token_id,\n",
    "                                        do_sample=True, top_p=0.7, temperature=0.95,)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 翻译能力几乎丧失\n",
    "response, history = base_model.chat(tokenizer, \"翻译：nice to see you\", history=[],max_length=512,\n",
    "                                        eos_token_id=config.eos_token_id,\n",
    "                                        do_sample=True, top_p=0.7, temperature=0.95,)\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
